#!/usr/bin/env python3
"""
MERIT-Basins 简化边界数据下载脚本
下载 catchments_simplified.zip 用于低分辨率模式

使用方法:
    python simplified_catchments.py
    nohup python simplified_catchments.py > ../logs/catchments.log 2>&1 &
"""

import os
import sys
import requests
from loguru import logger
from tqdm import tqdm

# 配置
URL = "https://mghydro.com/watersheds/share/catchments_simplified.zip"
BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "mghydro", "catchments")
FILENAME = "catchments_simplified.zip"

def setup_logger():
    """配置日志系统"""
    logger.remove()
    is_nohup = not sys.stdout.isatty()
    
    # 确保logs目录存在
    log_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "logs")
    os.makedirs(log_dir, exist_ok=True)
    log_file = os.path.join(log_dir, "catchments_download.log")
    
    if is_nohup:
        logger.add(log_file, rotation="10 MB", retention="7 days",
                  format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}", level="INFO")
        logger.add(sys.stdout, format="{time:HH:mm:ss} | {message}", level="INFO", colorize=False)
    else:
        logger.add(sys.stdout, format="<green>{time:HH:mm:ss}</green> | <level>{level}</level> | {message}", 
                  level="INFO", colorize=True)
        logger.add(log_file, rotation="10 MB", retention="7 days",
                  format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}", level="DEBUG")

def download_file():
    """下载catchments_simplified.zip文件"""
    local_path = os.path.join(BASE_DIR, FILENAME)
    
    # 创建目录
    os.makedirs(BASE_DIR, exist_ok=True)
    
    # 检查文件是否已存在
    if os.path.exists(local_path):
        logger.info(f"文件已存在: {local_path}")
        file_size = os.path.getsize(local_path)
        logger.info(f"文件大小: {file_size:,} 字节")
        return
    
    try:
        logger.info(f"开始下载: {URL}")
        response = requests.get(URL, stream=True, timeout=30)
        response.raise_for_status()
        
        total_size = int(response.headers.get('content-length', 0))
        logger.info(f"文件大小: {total_size:,} 字节")
        
        with open(local_path, 'wb') as f, tqdm(
            desc="下载进度",
            total=total_size,
            unit='B',
            unit_scale=True,
            unit_divisor=1024,
        ) as pbar:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    pbar.update(len(chunk))
        
        final_size = os.path.getsize(local_path)
        logger.success(f"下载完成: {local_path} ({final_size:,} 字节)")
        
        # 提示解压
        logger.info("下载完成后，请手动解压文件:")
        logger.info(f"cd {BASE_DIR}")
        logger.info(f"unzip {FILENAME}")
        logger.info("然后在 delineator 的 config.py 中设置:")
        logger.info(f"LOWRES_CATCHMENTS_DIR = '{BASE_DIR}/解压后的目录'")
        
    except Exception as e:
        logger.error(f"下载失败: {e}")
        if os.path.exists(local_path):
            os.remove(local_path)
        raise

def main():
    """主函数"""
    setup_logger()
    logger.info("开始下载 MERIT-Basins 简化边界数据")
    logger.info(f"目标目录: {BASE_DIR}")
    
    try:
        download_file()
        logger.success("任务完成!")
    except Exception as e:
        logger.error(f"任务失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()